import os
import sys
import torch
import torch.nn as nn
import numpy as np
parent_path = os.path.dirname(sys.path[0])
if parent_path not in sys.path:
    sys.path.append(parent_path)

from lib.resnet import resnet50, GCC
from lib.bninception import bninception
from lib.clustering import cluster_by_cap


def t_resnet():
    '''
    unit test for lib/resnet.py
    '''
    model = resnet50(pretrained=True).cuda()
    input = torch.rand(4, 3, 224, 224).cuda()
    x = model(input)
    assert (x.shape == torch.Size([4, 2048, 7, 7]))
    # print(x)
    print(x.shape)

    gcc_input = torch.rand(4, 128, 112, 112).cuda()
    gcc = GCC(128, 128, k=4).cuda()
    assert (gcc(gcc_input).shape == torch.Size([4, 128, 56, 56]))

    model.layer2[0].conv2 = gcc
    x = model(input)
    assert (x.shape == torch.Size([4, 2048, 7, 7]))
    __import__('ipdb').set_trace()


def t_resnet34():
    import torchvision.models as models
    model = models.resnet34(pretrained=True).cuda()
    input = torch.rand(4, 3, 224, 224).cuda()
    x = model(input)
    print(model)


from lib.loss.sampler import Sampler


def t_Sampler():
    '''
    test simple Sampler
    Sampler 先计算输入batch的距离矩阵，再计算相同类pos, 再计算不同类样本的neg
    '''
    s = Sampler(cutoff=0.5, infinity=1e6, eps=1e-6)
    x = np.array([[1.0, 1.0], [1.1, 1.1], [1.0, -1.0], [1.0, -1.1], [-1., -1.],
                  [-1., -1.1]])
    x = torch.tensor(x)
    labels = torch.tensor([1, 1, 2, 2, 3, 3])
    a_indices, a, p, n = s(x, labels)
    print(a_indices)
    print(a)
    print(p)
    print(n)


def t_bninception():
    model = bninception()


def t_inception_v3():
    from lib.inceptionv3 import inception_v3
    model = inception_v3().cuda()
    input = torch.rand(4, 3, 299, 299).cuda()
    o = model(input)


def t_cluster_by_cap():
    cap = np.array([
        [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8],
        [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8],
        [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8],
        [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8],
    ])
    # cap = torch.rand(4, 8054)
    rv = cluster_by_cap(cap)
    print(rv)


if __name__ == "__main__":
    t_resnet34()
    # t_cluster_by_cap()
    # t_resnet()
    # t_Sampler()
    # t_bninception()
    # t_inception_v3()
